import os
import shutil
from pathlib import Path

def batch_rename(image_dir, label_dir, start_num, copy_mode=False, output_image_dir=None, output_label_dir=None):
    # 验证输入文件夹是否存在
    if not os.path.exists(image_dir):
        print(f"错误：图片文件夹 '{image_dir}' 不存在")
        return
    
    if not os.path.exists(label_dir):
        print(f"错误：标签文件夹 '{label_dir}' 不存在")
        return
    
    # 如果使用复制模式，验证输出文件夹是否存在或创建
    if copy_mode:
        if output_image_dir:
            os.makedirs(output_image_dir, exist_ok=True)
        if output_label_dir:
            os.makedirs(output_label_dir, exist_ok=True)
    
    # 获取所有文件并按名称排序
    image_files = sorted(os.listdir(image_dir))
    label_files = sorted(os.listdir(label_dir))
    
    # 检查文件数量是否匹配
    if len(image_files) != len(label_files):
        print(f"警告：图片文件数量({len(image_files)})与标签文件数量({len(label_files)})不匹配")
    
    # 准备处理
    total = max(len(image_files), len(label_files))
    current_num = start_num
    
    print(f"开始处理 {total} 个文件对...")
    
    # 处理每个文件
    for i in range(total):
        # 获取当前文件
        img_file = image_files[i] if i < len(image_files) else None
        lbl_file = label_files[i] if i < len(label_files) else None
        
        # 生成新文件名
        new_base = f"frame_{current_num:05d}"
        
        # 处理图片文件
        if img_file:
            img_ext = Path(img_file).suffix
            new_img_name = new_base + img_ext
            
            source_img = os.path.join(image_dir, img_file)
            if copy_mode and output_image_dir:
                dest_img = os.path.join(output_image_dir, new_img_name)
                shutil.copy2(source_img, dest_img)
            else:
                dest_img = os.path.join(image_dir, new_img_name)
                os.rename(source_img, dest_img)
            
            print(f"图片: {img_file} -> {new_img_name}")
        
        # 处理标签文件
        if lbl_file:
            lbl_ext = Path(lbl_file).suffix
            new_lbl_name = new_base + lbl_ext
            
            source_lbl = os.path.join(label_dir, lbl_file)
            if copy_mode and output_label_dir:
                dest_lbl = os.path.join(output_label_dir, new_lbl_name)
                shutil.copy2(source_lbl, dest_lbl)
            else:
                dest_lbl = os.path.join(label_dir, new_lbl_name)
                os.rename(source_lbl, dest_lbl)
            
            print(f"标签: {lbl_file} -> {new_lbl_name}")
        
        current_num += 1
    
    print(f"处理完成。共重命名 {total} 个文件对。")

if __name__ == "__main__":
    # 直接在这里设置参数
    image_dir = r"D:\研究生\科研相关\无人艇专项\视觉数据集制作\total_202505211003\images"  # 图片文件夹路径
    label_dir = r"D:\研究生\科研相关\无人艇专项\视觉数据集制作\total_202505211003\annotations"  # 标签文件夹路径
    start_num = 00000       # 起始编号
    
    # 复制模式相关参数 (可选)
    copy_mode = False                # 是否使用复制模式
    output_image_dir = "./new_images"  # 输出图片文件夹 (仅在复制模式下使用)
    output_label_dir = "./new_labels"  # 输出标签文件夹 (仅在复制模式下使用)
    
    # 执行批量重命名
    batch_rename(image_dir, label_dir, start_num, copy_mode, output_image_dir, output_label_dir)    